import random
import string
import sys
import re

from WithSecure.common import fs

from drozer import android, util
from drozer.agent import builder, manifest
from drozer.server import uploader

class Exploit(object):
    
    def __init__(self, loader):
        self.__loader = loader
    
    def build_agent(self, arguments):
        source = arguments.no_gui and "rogue-agent" or "standard-agent"
        
        packager = builder.Packager()
        packager.unpack(source)
        
        if arguments.no_gui:
            e = manifest.Endpoint(packager.endpoint_path())
            if arguments.server != None:
                e.put_server(arguments.server)
            e.write()
            
        permissions = set(android.permissions)
        if arguments.permission != None:
            permissions = permissions.union(arguments.permission)
            
        m = manifest.Manifest(packager.manifest_path())
        for p in permissions:
            m.add_permission(p)
        m.write()
        
        return fs.read(packager.package())
    
    def build_multipart(self, parts, boundary):
        body = ""
        
        for k in parts:
            body += "--" + boundary + "; " + k + "\n"
            body += parts[k] + "\n"
        
        return body
    
    def execute(self, arguments):
        arguments.no_gui = True # there is no point in deploying a full GUI agent at this point...
        
        arguments.format = self.payload_format
        
        if len(self.payloads) > 0:
            if arguments.payload == None:
                sys.stderr.write("this exploit requires a payload\n")
                return
            if not arguments.payload in self.payloads:
                sys.stderr.write("unsupported payload: " + arguments.payload + "\n")
                return
        else:
            if arguments.payload != None:
                sys.stderr.write("skipping payload, this exploit does not require one\n")
                
                arguments.payload = None
        
        arguments.server = util.parse_server(arguments.server)
        if hasattr(self, "mode"):
            arguments.mode = self.mode
        if arguments.push_server != None:
            arguments.push_server = util.parse_server(arguments.push_server)
        # generate the payload
        if arguments.payload != None:
            self.payload = self.__loader.get(arguments.payload)(None, self.__loader, self).execute(arguments)
            # if the payload is for weasel, upload weasel to the server 
            if "weasel" in arguments.payload:
                if not self.upload_weasel(arguments):
                    return
        # generate the exploit
        self.generate(arguments)

    def generate_or_default_path(self, default):
        if default != None:
            path = default
            if not path.startswith("/"):
                path = "/" + path
        else:
            path = re.escape(self.generate_path())
        
        return path
            
    def generate_path(self, resource="view.jsp", token=True):
        if token:
            return "/%s?token=%s" % (resource, self.generate_token())
        else:
            return "/%s" % resource
    
    def generate_token(self, length=16):
        return ''.join(random.choice(string.ascii_uppercase + string.ascii_lowercase + string.digits) for x in range(length))
        
    def upload(self, arguments, resource, data, magic=None, mimetype=None, headers=None):
        if uploader.upload(arguments, resource, data, magic=magic, mimetype=mimetype, headers=headers):
            print(" [  OK  ]")
            
            return True
        else:
            print(" [ FAIL ]")
            
            return False
        
    def upload_weasel(self, arguments):
        weasel = self.__loader.get("weasel.shell.armeabi")(None, self.__loader, self)
        
        print("Uploading weasel to /weasel and W...")
        if not self.upload(arguments, "/weasel", weasel.execute(arguments), magic="W", mimetype="application/octet-stream"):
            return False
        print("Packaging an Agent... (this may take some time)")
        print("Uploading the Agent to /agent.apk and A...")
        if not self.upload(arguments, "/agent\\.apk", self.build_agent(arguments), magic="A", mimetype="application/vnd.android.package-archive"):
            return False
        
        return True
        
